from jax import random
from jax import numpy as np
from jax.nn import softmax
import rwkv_train_utils as tu


def get_tokenizer():
    # # opt 1
    # from data_utils import Tokenizer
    # tokenizer = Tokenizer.from_file("rwkv_tokenizer/20B_tokenizer.json")

    # opt 2
    from self_tokenizer import TRIE_TOKENIZER

    tokenizer = TRIE_TOKENIZER("rwkv_tokenizer/rwkv_vocab_v20230424.txt")

    # # opt 3
    # from rwkv_tokenizer import RWKV_TOKENIZER
    # tokenizer = RWKV_TOKENIZER('rwkv_tokenizer/rwkv_vocab_v20230424.txt')

    return tokenizer


def parse_rwkv_weight(pth_path="RWKV-4-Pile-169M-20220807-8023.pth"):
    import torch

    raw_weights = torch.load(pth_path, map_location="cpu")
    w = {}
    for key in raw_weights.keys():
        parts = key.split(".")
        last = parts.pop()
        current_ = w
        for p in parts:
            if p.isdigit():
                p = int(p)
            if p not in current_:
                current_[p] = {}
            current_ = current_[p]
        val = raw_weights[key]
        if ".time_" in key:
            val = val.squeeze()
        if ".time_decay" in key:
            val = -torch.exp(val)
        current_[last] = val.float().numpy()

    for i in w["blocks"].keys():
        att = w["blocks"][i]["att"]
        ffn = w["blocks"][i]["ffn"]

        for m in att, ffn:
            for k in ("key", "value", "receptance", "output"):
                if k in m:
                    m[k[0] + "_proj"] = m[k]["weight"]
                    del m[k]
    return w


def parse_model_info(weights_tree):
    return {
        "n_layer": len(weights_tree["blocks"]),
        "n_embed": weights_tree["emb"]["weight"].shape[1],
        "n_vocab": weights_tree["emb"]["weight"].shape[0],
    }


def rnn_generate(model, weights_tree, prompt, n_tokens=50, tokenizer=None, state=None):
    model_info = parse_model_info(weights_tree)

    if state is None:
        state = np.zeros((model_info["n_layer"], 5, model_info["n_embed"]))

        # this initialization solves the issue of dividing by 0 and the issue of not being able to stop
        for i in range(model_info["n_layer"]):
            state = state.at[i, 3].set(-1e30)
            state = state.at[i, 4].set(-1e30)
    
    if tokenizer is None:
        tokenizer = get_tokenizer()
    input_ids = tokenizer.encode(prompt)
    keygen = tu.KeyGen()

    for input_id in input_ids[:-1]:
        _, state = model(input_id, state, **weights_tree)
    input_id = input_ids[-1]

    import time

    start = time.time()
    n_gen = 0
    
    for _ in range(n_tokens):
        n_gen += 1
        logits, state = model(input_ids[-1], state, **weights_tree)
        out_id = top_p_sample(logits, keygen)
        # out_id = np.argmax(logits)
        if out_id < 2:
            break
        else:
            out_token = tokenizer.decode([int(out_id)])
        print(out_token, end="", flush=True)
        input_ids.append(out_id)

    end = time.time()
        
    print(f"\nSpeed: {(end-start)/n_gen:.2f} seconds / token\n")
    return state


def rnn_generate_batch_stateless(
    model, weights_tree, prompt, n_tokens=10, tokenizer=None
):
    if tokenizer is None:
        tokenizer = get_tokenizer()
    input_ids = tokenizer.encode(prompt)
    print(prompt, end="")
    keygen = tu.KeyGen()
    for _ in range(n_tokens):
        input_ids_batch = np.array(input_ids).reshape(1, -1)
        out = model(input_ids_batch, **weights_tree)
        # out_id = top_p_sample(out[0, -1, :], keygen)
        out_id = np.argmax(out[0, -1, :])
        res = tokenizer.decode([out_id])
        print(res, end="")
        input_ids.pop(0)  # jit becomes very slow when shape changes
        input_ids.append(out_id)
    print()


def top_p_threshold(probs, p):
    sorted_probs = np.sort(probs)[::-1]
    cumsum_probs = np.cumsum(softmax(sorted_probs), axis=-1)
    indices = np.where(cumsum_probs > p)[0]
    if len(indices) == 0:
        return 0.0
    else:
        return sorted_probs[indices[0]]


def top_p_sample(probs, keygen, temperature=0.2, top_p=0.9):
    threshold = top_p_threshold(probs, top_p)
    probs = np.where(probs < threshold, 0.0, probs)
    if temperature != 1.0:
        probs = np.power(probs, 1.0 / temperature)

    probs = probs / np.sum(probs)
    return random.choice(keygen(), a=len(probs), p=probs)
